import org.stringtemplate.v4.ST;
import org.stringtemplate.v4.STGroup;

import java.util.Objects;

public class NewCymbolFormatter extends CymbolBaseVisitor<ST> {

    private final STGroup templates;

    public NewCymbolFormatter(STGroup templates) {
        this.templates = templates;
    }

    @Override
    public ST visitCompilationUnit(CymbolParser.CompilationUnitContext ctx) {
        var fileTemplate = templates.getInstanceOf("file");

        for (var child : ctx.children) {
            fileTemplate.add("defs", visit(child));
        }

        return fileTemplate;
    }


    @Override
    public ST visitClassDeclaration(CymbolParser.ClassDeclarationContext ctx) {
        var classTemplate = templates.getInstanceOf("class");

        classTemplate.add("name", ctx.ID().getText());

        if (ctx.superClass() != null) {
            classTemplate.add("sup", ctx.superClass().ID().getText());
        }

        for (var member : ctx.classMember()) {
            classTemplate.add("members", visit(member));
        }

        return classTemplate;
    }

    @Override
    public ST visitNormalVarDeclaration(CymbolParser.NormalVarDeclarationContext ctx) {
        var varTemplate = templates.getInstanceOf("var");

        varTemplate.add("type", ctx.type().getText());
        varTemplate.add("name", ctx.ID().getText());
        varTemplate.add("ptr", false);

        if (ctx.expression() != null) {
            varTemplate.add("init", visit(ctx.expression()));
        }

        return varTemplate;
    }

    @Override
    public ST visitArrayVarDeclaration(CymbolParser.ArrayVarDeclarationContext ctx) {
        var varTemplate = templates.getInstanceOf("var");

        varTemplate.add("type", ctx.type().getText());
        varTemplate.add("name", ctx.ID().getText());
        varTemplate.add("ptr", true);

        if (ctx.expression() != null) {
            varTemplate.add("init", visit(ctx.expression()));
        }

        return varTemplate;
    }

    @Override
    public ST visitPointerVarDeclaration(CymbolParser.PointerVarDeclarationContext ctx) {
        var varTemplate = templates.getInstanceOf("var");

        varTemplate.add("type", ctx.type().getText());
        varTemplate.add("name", ctx.ID().getText());
        varTemplate.add("ptr", true);

        if (ctx.expression() != null) {
            varTemplate.add("init", visit(ctx.expression()));
        }

        return varTemplate;
    }

    @Override
    public ST visitMethodDeclaration(CymbolParser.MethodDeclarationContext ctx) {
        var methodTemplate = templates.getInstanceOf("method");

        methodTemplate.add("retType", ctx.type().getText());
        methodTemplate.add("name", ctx.ID().getText());

        if (ctx.formalParameters() != null) {
            for (var param : ctx.formalParameters().parameter()) {
                methodTemplate.add("args", visit(param));
            }
        }

        methodTemplate.add("block", visit(ctx.block()));
        return methodTemplate;
    }

    @Override
    public ST visitNormalParameter(CymbolParser.NormalParameterContext ctx) {
        var argTemplate = templates.getInstanceOf("arg");
        argTemplate.add("type", ctx.type().getText());
        argTemplate.add("name", ctx.ID().getText());
        argTemplate.add("ptr", false);
        return argTemplate;
    }

    @Override
    public ST visitArrayParameter(CymbolParser.ArrayParameterContext ctx) {
        var argTemplate = templates.getInstanceOf("arg");
        argTemplate.add("type", ctx.type().getText());
        argTemplate.add("name", ctx.ID().getText());
        argTemplate.add("ptr", true);
        return argTemplate;
    }

    @Override
    public ST visitPointerParameter(CymbolParser.PointerParameterContext ctx) {
        var argTemplate = templates.getInstanceOf("arg");
        argTemplate.add("type", ctx.type().getText());
        argTemplate.add("name", ctx.ID().getText());
        argTemplate.add("ptr", true);
        return argTemplate;
    }

    @Override
    public ST visitBlock(CymbolParser.BlockContext ctx) {
        var blockTemplate = templates.getInstanceOf("block");

        if (ctx.statement() != null) {
            for (var statement : ctx.statement()) {
                blockTemplate.add("stats", visit(statement));
            }
        }

        return blockTemplate;
    }

    @Override
    public ST visitIfStatement(CymbolParser.IfStatementContext ctx) {
        var ifTemplate = templates.getInstanceOf("if");

        ifTemplate.add("cond", visit(ctx.expression()));
        ifTemplate.add("stat1", visit(ctx.s));

        if (ctx.e != null) {
            ifTemplate.add("stat2", visit(ctx.e));
        }

        return ifTemplate;
    }

    @Override
    public ST visitReturnStatement(CymbolParser.ReturnStatementContext ctx) {
        var returnTemplate = templates.getInstanceOf("return");
        if (ctx.expression() != null) {
            returnTemplate.add("v", visit(ctx.expression()));
        }

        return returnTemplate;
    }

    @Override
    public ST visitAssignmentStatement(CymbolParser.AssignmentStatementContext ctx) {
        var assignTemplate = templates.getInstanceOf("assign");
        assignTemplate.add("a", visit(ctx.lhs()));
        assignTemplate.add("b", visit(ctx.expression()));
        return assignTemplate;
    }

    @Override
    public ST visitPostfixExpressionStatement(CymbolParser.PostfixExpressionStatementContext ctx) {
        var callStatTemplate = templates.getInstanceOf("callstat");

        var inner = visit(ctx.postfixExpression());
        callStatTemplate.add("name", inner.getAttribute("name"));
        callStatTemplate.add("args", inner.getAttribute("args"));

        return callStatTemplate;
    }

    @Override
    public ST visitPostfixExpression(CymbolParser.PostfixExpressionContext ctx) {
        if (ctx.primary() != null) {
            return visit(ctx.primary());
        }

        return visit(ctx.postfix()); // 在各个 postfix 的子方法里再去 visit(postfixExpression)
    }

    @Override
    public ST visitFunctionCallPostfix(CymbolParser.FunctionCallPostfixContext ctx) {
        var callTemplate = templates.getInstanceOf("call");
        var primary = ((CymbolParser.PostfixExpressionContext) (ctx.getParent())).postfixExpression();
        callTemplate.add("name", visit(primary));
        if (ctx.expressionList() != null) {
            for (var expr : ctx.expressionList().expr()) {
                callTemplate.add("args", visit(expr));
            }
        }

        return callTemplate;
    }

    @Override
    public ST visitArrayIndexPostfix(CymbolParser.ArrayIndexPostfixContext ctx) {
        var indexTemplate = templates.getInstanceOf("index");
        var primary = ((CymbolParser.PostfixExpressionContext) (ctx.getParent())).postfixExpression();
        indexTemplate.add("array", visit(primary));
        indexTemplate.add("i", visit(ctx.expr()));
        return indexTemplate;
    }

    @Override
    public ST visitMemberAccessPostfix(CymbolParser.MemberAccessPostfixContext ctx) {
        var memberTemplate = templates.getInstanceOf("member");
        var primary = ((CymbolParser.PostfixExpressionContext) (ctx.getParent())).postfixExpression();
        memberTemplate.add("obj", visit(primary));
        memberTemplate.add("name", ctx.ID().getText());
        return memberTemplate;
    }

    @Override
    public ST visitPointerMemberAccessPostfix(CymbolParser.PointerMemberAccessPostfixContext ctx) {
        // 转换为 (*p).x
        var dereferenceTemplate = templates.getInstanceOf("deref");
        var primary = ((CymbolParser.PostfixExpressionContext) (ctx.getParent())).postfixExpression();
        dereferenceTemplate.add("v", visit(primary));

        var memberTemplate = templates.getInstanceOf("member");
        memberTemplate.add("obj", dereferenceTemplate);
        memberTemplate.add("name", ctx.ID().getText());
        return memberTemplate;
    }

    @Override
    public ST visitExpr(CymbolParser.ExprContext ctx) {
        if (ctx.unaryExpression() != null) {
            return visit(ctx.unaryExpression());
        }

        var operationTemplate = templates.getInstanceOf("operation");
        operationTemplate.add("op", ctx.op.getText());
        operationTemplate.add("x", visit(ctx.expr(0)));
        operationTemplate.add("y", visit(ctx.expr(1)));
        return operationTemplate;
    }

    @Override
    public ST visitUnaryMinusExpression(CymbolParser.UnaryMinusExpressionContext ctx) {
        var unaryMinusTemplate = templates.getInstanceOf("unary_minus");
        unaryMinusTemplate.add("v", visit(ctx.unaryExpression()));
        return unaryMinusTemplate;
    }

    @Override
    public ST visitUnaryNotExpression(CymbolParser.UnaryNotExpressionContext ctx) {
        var unaryNotTemplate = templates.getInstanceOf("unary_not");
        unaryNotTemplate.add("v", visit(ctx.unaryExpression()));
        return unaryNotTemplate;
    }

    @Override
    public ST visitUnaryAddressExpression(CymbolParser.UnaryAddressExpressionContext ctx) {
        var addressTemplate = templates.getInstanceOf("addr");
        addressTemplate.add("v", visit(ctx.unaryExpression()));
        return addressTemplate;
    }

    @Override
    public ST visitUnaryArrayIndexExpression(CymbolParser.UnaryArrayIndexExpressionContext ctx) {
        // 处理 *(p+i) 的情况，此时要转换成 p[i]
        var indexTemplate = templates.getInstanceOf("index");
        indexTemplate.add("array", ctx.ID().getText());
        indexTemplate.add("i", visit(ctx.expr()));
        return indexTemplate;
    }

    @Override
    public ST visitUnaryDereferenceExpression(CymbolParser.UnaryDereferenceExpressionContext ctx) {
        var dereferenceTemplate = templates.getInstanceOf("deref");
        dereferenceTemplate.add("v", visit(ctx.unaryExpression()));
        return dereferenceTemplate;
    }

    @Override
    public ST visitPostfixExpressionUnary(CymbolParser.PostfixExpressionUnaryContext ctx) {
        return visit(ctx.postfixExpression());
    }

    @Override
    public ST visitPrimary(CymbolParser.PrimaryContext ctx) {
        if (ctx.expr() != null) {
            return visit(ctx.expr());
        }

        var operatorTemplate = templates.getInstanceOf("operator");
        operatorTemplate.add("o", ctx.getText());
        return operatorTemplate;
    }
}
